package org.springframework.security.oauth2.client.login.ext;

import org.springframework.security.oauth2.client.OAuth2ClientContext;
import org.springframework.security.oauth2.client.OAuth2RestTemplate;
import org.springframework.security.oauth2.client.resource.OAuth2ProtectedResourceDetails;
import org.springframework.security.oauth2.client.resource.UserRedirectRequiredException;
import org.springframework.security.oauth2.client.token.AccessTokenRequest;
import org.springframework.security.oauth2.common.OAuth2AccessToken;

public class OAuth2RestStateTemplate extends OAuth2RestTemplate {

	public OAuth2RestStateTemplate(OAuth2ProtectedResourceDetails resource) {
		super(resource);
	}

	public OAuth2RestStateTemplate(OAuth2ProtectedResourceDetails resource, OAuth2ClientContext context) {
		super(resource, context);
	}

	protected OAuth2AccessToken acquireAccessToken(OAuth2ClientContext oauth2Context)
			throws UserRedirectRequiredException {
		AccessTokenRequest accessTokenRequest = oauth2Context.getAccessTokenRequest();
		String stateKey = accessTokenRequest.getStateKey();
		if (stateKey != null) {
			oauth2Context.setPreservedState(stateKey, "http://127.0.0.1:8081/login");
		}
		return super.acquireAccessToken(oauth2Context);
	}

}
